import or1200_definitions
import sys
sys.path.append("..")
import pyrtl
from pyrtl import  GPUSim_now
width = or1200_definitions.OR1200_OPERAND_WIDTH

class Or1200Alu(object):
    def __init__(self):
        self.a = pyrtl.Input(bitwidth=width, name='a')
        self.b = pyrtl.Input(bitwidth=width, name='b')
        self.alu_op = pyrtl.Input(bitwidth=or1200_definitions.OR1200_ALUOP_WIDTH, name='alu_op')
        self.shrot_op = pyrtl.Input(bitwidth=or1200_definitions.OR1200_SHROTOP_WIDTH, name='shrot_op')
        self.comp_op = pyrtl.Input(bitwidth=or1200_definitions.OR1200_COMPOP_WIDTH, name='comp_op')
        self.carry = pyrtl.Input(bitwidth=1, name='carry')
        self.flag = pyrtl.Input(bitwidth=1, name='flag')

        self.result = pyrtl.Output(bitwidth=width, name='result')
        self.flagforw = pyrtl.Output(bitwidth=1, name='flagforw')
        self.flag_we = pyrtl.Output(bitwidth=1, name='flag_we')
        self.cyforw = pyrtl.Output(bitwidth=1, name='cyforw')
        self.cy_we = pyrtl.Output(bitwidth=1, name='cy_we')

        self.comp_a = pyrtl.WireVector(bitwidth=width)
        self.comp_b = pyrtl.WireVector(bitwidth=width)
        self.result_sum = pyrtl.WireVector(bitwidth=width)
        self.cy_sum = pyrtl.WireVector(bitwidth=1)
        self.result_csum = pyrtl.WireVector(bitwidth=width)
        self.cy_csum = pyrtl.WireVector(bitwidth=1)

        self.shifted_rotated = pyrtl.WireVector(bitwidth=width)
        self.result_and = pyrtl.WireVector(bitwidth=width)
        self.flagcomp = pyrtl.WireVector(bitwidth=1)

        #initialize all the parts
        general_alu = GeneralAlu()
        generate_flag = GenerateFlag()
        generate_sry = GenerateSRY()
        shift_rotate = ShiftAndRotate()
        compare_impl = CompareImpl()

        self.result <<= general_alu.result
        general_alu.a <<= self.a
        general_alu.b <<= self.b
        general_alu.alu_op <<= self.alu_op
        general_alu.result_sum <<= self.result_sum
        general_alu.result_csum <<= self.result_csum
        general_alu.result_and <<= self.result_and
        general_alu.shifted_rotated <<= self.shifted_rotated
        general_alu.flag <<= self.flag
        generate_flag.alu_op <<= self.alu_op
        generate_flag.result_sum <<= self.result_sum
        generate_flag.flagcomp <<= self.flagcomp
        self.flagforw <<= generate_flag.flagforw
        self.flag_we <<= generate_flag.flag_we
        generate_flag.result_and <<= self.result_and
        generate_sry.alu_op <<= self.alu_op
        generate_sry.cy_sum <<= self.cy_sum
        generate_sry.cy_csum <<= self.cy_csum
        self.cy_we <<= generate_sry.cy_we
        self.cyforw <<= generate_sry.cyforw
        shift_rotate.a <<= self.a
        shift_rotate.b <<= self.b
        shift_rotate.shrot_op <<= self.shrot_op
        self.shifted_rotated <<= shift_rotate.shifted_rotated
        compare_impl.comp_a <<= self.comp_a
        compare_impl.comp_b <<= self.comp_b
        compare_impl.comp_op <<= self.comp_op
        self.flagcomp <<= compare_impl.flagcomp

        self.comp_a <<= pyrtl.concat(self.a[-1] ^ self.comp_op[3], self.a[:-2])
        self.comp_b <<= pyrtl.concat(self.b[-1] ^ self.comp_op[3], self.b[:-2])
        temp = self.a + self.b
        self.cy_sum <<= temp[-1]
        self.result_sum <<= temp[:-2]
        if or1200_definitions.OR1200_IMPL_ADDC:
            temp2 = self.a + self.b + pyrtl.concat(pyrtl.Const(0, bitwidth=32), self.carry)
            self.cy_csum <<= temp2[-1]
            self.result_csum <<= temp2[:-2]

        self.result_and <<= self.a & self.b

class Or1200Alu_wire(object):
    def __init__(self):
        self.a = pyrtl.WireVector(bitwidth=width, name='alu_a')
        self.b = pyrtl.WireVector(bitwidth=width, name='alu_b')
        self.alu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_ALUOP_WIDTH, name='alu_alu_op')
        self.shrot_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_SHROTOP_WIDTH, name='alu_shrot_op')
        self.comp_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_COMPOP_WIDTH, name='alu_comp_op')
        self.carry = pyrtl.WireVector(bitwidth=1, name='alu_carry')
        self.flag = pyrtl.WireVector(bitwidth=1, name='alu_flag')

        self.result = pyrtl.WireVector(bitwidth=width, name='alu_result')
        self.flagforw = pyrtl.WireVector(bitwidth=1, name='alu_flagforw')
        self.flag_we = pyrtl.WireVector(bitwidth=1, name='alu_flag_we')
        self.cyforw = pyrtl.WireVector(bitwidth=1, name='alu_cyforw')
        self.cy_we = pyrtl.WireVector(bitwidth=1, name='alu_cy_we')

        self.comp_a = pyrtl.WireVector(bitwidth=width)
        self.comp_b = pyrtl.WireVector(bitwidth=width)
        self.result_sum = pyrtl.WireVector(bitwidth=width)
        self.cy_sum = pyrtl.WireVector(bitwidth=1)
        self.result_csum = pyrtl.WireVector(bitwidth=width)
        self.cy_csum = pyrtl.WireVector(bitwidth=1)

        self.shifted_rotated = pyrtl.WireVector(bitwidth=width)
        self.result_and = pyrtl.WireVector(bitwidth=width)
        self.flagcomp = pyrtl.WireVector(bitwidth=1)

        # initialize all the parts
        general_alu = GeneralAlu()
        generate_flag = GenerateFlag()
        generate_sry = GenerateSRY()
        shift_rotate = ShiftAndRotate()
        compare_impl = CompareImpl()

        self.result <<= general_alu.result
        general_alu.a <<= self.a
        general_alu.b <<= self.b
        general_alu.alu_op <<= self.alu_op
        general_alu.result_sum <<= self.result_sum
        general_alu.result_csum <<= self.result_csum
        general_alu.result_and <<= self.result_and
        general_alu.shifted_rotated <<= self.shifted_rotated
        general_alu.flag <<= self.flag
        generate_flag.alu_op <<= self.alu_op
        generate_flag.result_sum <<= self.result_sum
        generate_flag.flagcomp <<= self.flagcomp
        self.flagforw <<= generate_flag.flagforw
        self.flag_we <<= generate_flag.flag_we
        generate_flag.result_and <<= self.result_and
        generate_sry.alu_op <<= self.alu_op
        generate_sry.cy_sum <<= self.cy_sum
        generate_sry.cy_csum <<= self.cy_csum
        self.cy_we <<= generate_sry.cy_we
        self.cyforw <<= generate_sry.cyforw
        shift_rotate.a <<= self.a
        shift_rotate.b <<= self.b
        shift_rotate.shrot_op <<= self.shrot_op
        self.shifted_rotated <<= shift_rotate.shifted_rotated
        compare_impl.comp_a <<= self.comp_a
        compare_impl.comp_b <<= self.comp_b
        compare_impl.comp_op <<= self.comp_op
        self.flagcomp <<= compare_impl.flagcomp

        self.comp_a <<= pyrtl.concat(self.a[-1] ^ self.comp_op[3], self.a[:-2])
        self.comp_b <<= pyrtl.concat(self.b[-1] ^ self.comp_op[3], self.b[:-2])
        temp = self.a + self.b
        self.cy_sum <<= temp[-1]
        self.result_sum <<= temp[:-2]
        if or1200_definitions.OR1200_IMPL_ADDC:
            temp2 = self.a + self.b + pyrtl.concat(pyrtl.Const(0, bitwidth=32), self.carry)
            self.cy_csum <<= temp2[-1]
            self.result_csum <<= temp2[:-2]

        self.result_and <<= self.a & self.b

class GeneralAlu(object):
    def __init__(self):
        self.result = pyrtl.WireVector(bitwidth=width)
        self.a = pyrtl.WireVector(bitwidth=width)
        self.b = pyrtl.WireVector(bitwidth=width)
        self.alu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_ALUOP_WIDTH)
        self.result_sum = pyrtl.WireVector(bitwidth=width)
        self.result_csum = pyrtl.WireVector(bitwidth=width)
        self.result_and = pyrtl.WireVector(bitwidth=width)
        self.shifted_rotated = pyrtl.WireVector(bitwidth=width)
        self.flag = pyrtl.WireVector(bitwidth=1)

        with pyrtl.conditional_assignment:
            with self.alu_op == or1200_definitions.OR1200_ALUOP_FF1:
                # tmp = pyrtl.select(self.a[0] == or1200_definitions.ONE, 1, 0)
                tmp = pyrtl.select(self.a[0].__and__(or1200_definitions.ONE), 1, 0)
                for i in range(1, width):
                    tmp = pyrtl.select(self.a[i], i+1, tmp)
                self.result |= tmp
            with self.alu_op == or1200_definitions.OR1200_ALUOP_SHROT:
                self.result |= self.shifted_rotated
            with self.alu_op == or1200_definitions.OR1200_ALUOP_ADD:
                self.result |= self.result_sum
            with self.alu_op == or1200_definitions.OR1200_ALUOP_ADDC:
                self.result |= self.result_csum
            with self.alu_op == or1200_definitions.OR1200_ALUOP_SUB:
                self.result |= self.a - self.b
            with self.alu_op == or1200_definitions.OR1200_ALUOP_XOR:
                self.result |= self.a ^ self.b
            with self.alu_op == or1200_definitions.OR1200_ALUOP_OR:
                self.result |= self.a | self.b
            with self.alu_op == or1200_definitions.OR1200_ALUOP_IMM:
                self.result |= self.b
            with self.alu_op == or1200_definitions.OR1200_ALUOP_MOVHI:
                self.result |= pyrtl.shift_left_logical(self.b, pyrtl.Const(16))
            with self.alu_op == or1200_definitions.OR1200_ALUOP_CMOV:
                self.result |= pyrtl.select(self.flag, self.a, self.b)
            with self.alu_op == or1200_definitions.OR1200_ALUOP_COMP:
                self.result |= self.result_and
            with self.alu_op == or1200_definitions.OR1200_ALUOP_AND:
                self.result |= self.result_and

class GenerateFlag(object):
    def __init__(self):
        self.alu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_ALUOP_WIDTH)
        self.result_sum = pyrtl.WireVector(bitwidth=width)
        self.flagcomp = pyrtl.WireVector(bitwidth=1)
        self.flagforw = pyrtl.WireVector(bitwidth=1)
        self.flag_we = pyrtl.WireVector(bitwidth=1) #flag write enable
        self.result_and = pyrtl.WireVector(bitwidth=width)

        with pyrtl.conditional_assignment:
            with self.alu_op == or1200_definitions.OR1200_ALUOP_COMP:
                self.flagforw |= self.flagcomp
                self.flag_we |= pyrtl.Const(0b1, bitwidth=1)
            with pyrtl.otherwise:
                self.flagforw |= pyrtl.Const(0b0, bitwidth=1)
                self.flag_we |= pyrtl.Const(0b0, bitwidth=1)

class GenerateSRY(object):
    def __init__(self):
        self.alu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_ALUOP_WIDTH)
        self.cy_sum = pyrtl.WireVector(bitwidth=1)
        self.cy_csum = pyrtl.WireVector(bitwidth=1)
        self.cy_we = pyrtl.WireVector(bitwidth=1) #cy write enable
        self.cyforw = pyrtl.WireVector(bitwidth=1)

        with pyrtl.conditional_assignment:
            with self.alu_op == or1200_definitions.OR1200_ALUOP_ADD:
                self.cyforw |= self.cy_sum
                self.cy_we |= pyrtl.Const(0b1, bitwidth=1)
            with self.alu_op == or1200_definitions.OR1200_ALUOP_ADDC:
                self.cyforw |= self.cy_csum
                self.cy_we |= pyrtl.Const(0b1, bitwidth=1)
            with pyrtl.otherwise:
                self.cyforw |= pyrtl.Const(0b0, bitwidth=1)
                self.cy_we |= pyrtl.Const(0b0, bitwidth=1)

class ShiftAndRotate(object):
    def __init__(self):
        self.a = pyrtl.WireVector(bitwidth=width)
        self.b = pyrtl.WireVector(bitwidth=width)
        self.shrot_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_SHROTOP_WIDTH)
        self.shifted_rotated = pyrtl.WireVector(bitwidth=width)

        with pyrtl.conditional_assignment:
            with self.shrot_op == or1200_definitions.OR1200_SHROTOP_SLL:
                self.shifted_rotated |= pyrtl.shift_left_logical(self.a, self.b[0:5])
            with self.shrot_op == or1200_definitions.OR1200_SHROTOP_SRL:
                self.shifted_rotated |= pyrtl.shift_right_logical(self.a, self.b[0:5])
            with pyrtl.otherwise:
                #{32{a[31]}} << (6'd32-{1'b0, b[4:0]})) | a >> b[4:0]
                self.shifted_rotated |= (pyrtl.shift_left_logical(pyrtl.concat_list([self.a[-1]]*32), pyrtl.Const(32, bitwidth=6)) - pyrtl.concat(pyrtl.Const(0b0, bitwidth=1), self.b[0:5])) \
                                        | (pyrtl.shift_right_logical(self.a, self.b[0:5]))

class CompareImpl(object):
    def __init__(self):
        self.comp_a = pyrtl.WireVector(bitwidth=width)
        self.comp_b = pyrtl.WireVector(bitwidth=width)
        self.comp_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_COMPOP_WIDTH)
        self.flagcomp = pyrtl.WireVector(bitwidth=1)

        with pyrtl.conditional_assignment:
            with self.comp_op[0:3] == or1200_definitions.OR1200_COP_SFEQ:
                self.flagcomp |= (self.comp_a == self.comp_b)
            with self.comp_op[0:3] == or1200_definitions.OR1200_COP_SFNE:
                self.flagcomp |= (self.comp_a != self.comp_b)
            with self.comp_op[0:3] == or1200_definitions.OR1200_COP_SFGT:
                self.flagcomp |= (self.comp_a > self.comp_b)
            with self.comp_op[0:3] == or1200_definitions.OR1200_COP_SFGE:
                self.flagcomp |= (self.comp_a >= self.comp_b)
            with self.comp_op[0:3] == or1200_definitions.OR1200_COP_SFLT:
                self.flagcomp |= (self.comp_a < self.comp_b)
            with self.comp_op[0:3] == or1200_definitions.OR1200_COP_SFLE:
                self.flagcomp |= (self.comp_a <= self.comp_b)
            with pyrtl.otherwise:
                self.flagcomp |= pyrtl.Const(0b0, bitwidth=1)

if __name__ == '__main__':
    alu = Or1200Alu()
    sim = pyrtl.GPUSim_now.GPUSim('alu')

    com = pyrtl.CompiledSimulation()

    input_width = []
    from pyrtl.wire import Input

    IN = pyrtl.working_block().wirevector_subset(Input)
    for inn in IN:
        name = inn.name
        width = inn.bitwidth
        input_width.append([name, width])
    import random

    max_ = 0
    com_time = 0
    ess_time = 0
    num = 1
    for i in range(num):
        print(i)

        inputs_random = {}
        inputs_select = {}
        inputs_old = {}
        inputs_ = {}

        for inp in input_width:
            inputs_random[inp[0]] = []
            inputs_select[inp[0]] = []
            result = random.randint(0, 2 ** inp[1] - 1)
            inputs_old[inp[0]] = result

        for j in range(131072):
            for inp in input_width:
                result = random.randint(0, 2 ** inp[1] - 1)
                inputs_[inp[0]] = result
                inputs_random[inp[0]].append(result)
            if random.random() < 0.2:
                for name in inputs_old:
                    inputs_select[name].append(inputs_old[name])
            else:
                for name in inputs_:
                    inputs_select[name].append(inputs_[name])
                    inputs_old[name] = inputs_[name]
        time1 = com.step_multiple(inputs_select)

    sim.make_input(inputs_select)